#include "precompiled.h"
#include "IMLGenerator.h"
#include "IMLEmitter.h"

namespace IcyCrystal {
	IMLGenerator::IMLGenerator(Namespace& globalNamespace, SyntaxTreeList& syntaxTree) : globalNamespace(globalNamespace), syntaxTree(syntaxTree) {
		// Create the registers!
		registerArray = new Register[5];
		registerArray[0].registerName = "EAX";
		registerArray[0].variables = new SymbolList;
		registerArray[1].registerName = "ECX";
		registerArray[1].variables = new SymbolList;
		registerArray[2].registerName = "EDX";
		registerArray[2].variables = new SymbolList;
		registerArray[3].registerName = "EDI";
		registerArray[3].variables = new SymbolList;
		registerArray[4].registerName = "ESI";
		registerArray[4].variables = new SymbolList;

		temporaryCount = 0;
	}

	GlobalDeclaration IMLGenerator::createArrayDeclaration(Symbol* sym){
		// Simply linearly output the values.
		GlobalDeclaration ret;
		ret.globalType = ARRAY;
		ret.symbolName = &sym->symbolName;
		ret.value.arrayValue = new vector<GlobalDeclaration>;

		if( sym->type.primitiveType == PRIM_TYPE_ARRAY ){
			// Add all values to the array.
			SyntaxTreeList::iterator it;
			for(it = sym->consts.exprConst->extra.arrayInitArguments->begin(); it!= sym->consts.exprConst->extra.arrayInitArguments->end(); it++){
				vector<GlobalDeclaration>* vec;
				vec = constToDeclaration(*it);
				ret.value.arrayValue->insert(ret.value.arrayValue->end(), vec->begin(), vec->end());
				delete vec;
			}
		}
		else if( sym->type.primitiveType == PRIM_TYPE_STRING ){
			ret.globalType = ARRAY;
			ret.symbolName = &sym->symbolName;
			GlobalDeclaration theString;
			theString.globalType = STRING;
			theString.value.stringValue = sym->consts.stringConst;
			ret.value.arrayValue->push_back(theString);
		}

		return ret;
	}

	vector<GlobalDeclaration>* IMLGenerator::constToDeclaration(SyntaxTree* tree){
		vector<GlobalDeclaration>* ret = new vector<GlobalDeclaration>;
		GlobalDeclaration decl;
		switch(tree->nodeType){
		case NODE_TOKEN:
			switch(tree->expressionReturnType->primitiveType){
			case PRIM_TYPE_INT:
				decl.globalType = INT;
				decl.value.intValue = tree->inner.nodeToken.value.intValue;
				ret->push_back(decl);
				break;
			case PRIM_TYPE_FLOAT:
				decl.globalType = FLOAT;
				decl.value.floatValue = tree->inner.nodeToken.value.floatValue;
				ret->push_back(decl);
				break;
			}
			break;
		case NODE_SYMBOL:
			if( tree->nodeFlags == NODE_IS_ARRAY_INITIATOR ){
				// Add
				SyntaxTreeList::iterator it;
				for( it = tree->extra.arrayInitArguments->begin(); it != tree->extra.arrayInitArguments->end(); it++){
					vector<GlobalDeclaration>* args = constToDeclaration(*it);
					ret->insert(ret->end(), args->begin(), args->end());
					delete args;
				}
			}
			break;
		default:
			throw new exception();
		}

		return ret;
	}

	Arg IMLGenerator::getNewTemporary(){
		stringstream s;
		s << "t" << temporaryCount;
		Arg ret;
		ret.arg = TEMP;
		ret.name.symbolName = new string(s.str());
		temporaryCount++;

		return ret;
	}
	void IMLGenerator::freeLastTemporary(){
		temporaryCount--;
	}

	ArgType IMLGenerator::scopeToArgType(Symbol* sym){
		switch(sym->scope){
		case SCOPE_LOCAL:
			return LOCAL;
		case SCOPE_GLOBAL:
			return GLOBAL;
		case SCOPE_TEMP:
			return TEMP;
		case SCOPE_CONSTANT:
			return CONSTANT;
		}
	}

	ArgList* IMLGenerator::generateExpression( Namespace* currentNamespace, SyntaxTree* expr, Arg& finalDest){
		/**
		* First pass expression generation.
		* This is the higest level IML.
		*/
		ArgList* ret = new ArgList;
		Symbol* sym;
		Quad* st;
		Token* tok;
		
		switch( expr->nodeType ){

		case NODE_SYMBOL:
			// Store the symbol into a temporary
			//st = new Quad;
			//st->op = MOV;
			sym = expr->inner.sym;
			if( sym->scope == SCOPE_GLOBAL ){
				finalDest.arg = GLOBAL;
				finalDest.name.symbolName = &sym->symbolName;
			}
			else if( sym->scope == SCOPE_LOCAL || sym->scope == SCOPE_PARAM){
				// Since this is still the first pass, generate a label.
				finalDest.arg = LOCAL;
				finalDest.name.symbolName = &sym->symbolName;
			}
			else throw new exception();
			
			delete ret;

			return NULL;
			break;

		case NODE_TOKEN:
			// Store the constant into a temporary
			st = new Quad;
			tok = & expr->inner.nodeToken;
			if( tok->tokenType == TOKEN_INT){
				finalDest.arg = INT_CONSTANT;
				finalDest.name.intName = tok->value.intValue;
			}
			else if( tok->tokenType == TOKEN_FLOAT ){
				finalDest.arg = FLOAT_CONSTANT;
				throw new exception(); // TODO: Handle this.
			}

			st->op = MOV;
			st->inner.normalOp.arg1 = finalDest;
			st->inner.normalOp.result = getNewTemporary();

			ret->push_back(st);
			finalDest = st->inner.normalOp.result;
			break;

		case NODE_EXPRESSION:
			// Our destination is a temporary.
			Arg arg1, arg2;
			tok = &expr->inner.nodeToken;
			if( tok->tokenType == TOKEN_PLUS || tok->tokenType == TOKEN_MINUS || tok->tokenType == TOKEN_MULT || tok->tokenType == TOKEN_DIV ){
				// Operation that requires two arguments.
				
				ArgList* lvalue = generateExpression(currentNamespace, expr->left, arg1);
				if( lvalue != NULL ){
					// Add the computation.
					ret->insert(ret->end(), lvalue->begin(), lvalue->end());
					delete lvalue;
				}
				
				ArgList* rvalue = generateExpression(currentNamespace, expr->right, arg2);
				if( rvalue != NULL ){
					ret->insert(ret->end(), rvalue->begin(), rvalue->end());
					delete rvalue;
				}

				// Generate our expression.
				Quad* quad = new Quad;
				switch(tok->tokenType){
				case TOKEN_PLUS:
					quad->op = ADD;
					break;
				case TOKEN_MINUS:
					quad->op = SUB;
					break;
				case TOKEN_MULT:
					quad->op = MULT;
					break;
				case TOKEN_DIV:
					quad->op = DIV;
					break;
				}
				quad->inner.normalOp.arg1 = arg1;
				quad->inner.normalOp.arg2 = arg2;
				quad->inner.normalOp.result = getNewTemporary();
				finalDest = quad->inner.normalOp.result;

				ret->push_back(quad);
				// Free the temporaries.
				if( arg1.arg == TEMP )
					freeLastTemporary();
				
			}
			else if( tok->tokenType == TOKEN_CONV_LITERAL ){
				delete ret;

				// Take this symbol literally.

				// Find the root node.
				while( expr->left != NULL )
					expr = expr->left;
				finalDest.arg = scopeToArgType(expr->inner.sym);
				finalDest.name.symbolName = &expr->inner.sym->symbolName;
				return NULL;
			}
			break;
		}

		return ret;

		
	}

	ArgList* IMLGenerator::generateStatement(Namespace* currentNamespace, SyntaxTree* statement){
		ArgList* ret = new ArgList;
		if( statement->nodeType == NODE_ASSIGN ){
			Arg finalDest;
			ArgList* genExpr = generateExpression(currentNamespace,statement->right, finalDest);

			if( genExpr != NULL){
				ret->insert(ret->end(), genExpr->begin(), genExpr->end());
				delete genExpr;
			}

			// Store the final destination into us.
			Quad* st = new Quad;
			st->op = MOV;
			st->inner.normalOp.arg1 = finalDest;
			Arg theSymbol;
			theSymbol.arg = scopeToArgType(statement->left->inner.sym);
			theSymbol.name.symbolName = &statement->left->inner.sym->symbolName;
			st->inner.normalOp.result = theSymbol;

			ret->push_back(st);

			//if( finalDest.arg == TEMP )
			//	freeLastTemporary();
		}

		return ret;
	}

	ArgList* IMLGenerator::generateFunctionBody(Symbol* funcSymbol){
		ArgList* ret = new ArgList;
		SyntaxTreeList* statements =funcSymbol->consts.function.functionBody;
		SyntaxTreeList::iterator it;

		for( it = statements->begin(); it!= statements->end(); it++){
			SyntaxTree* tree = *it;

			ArgList* stat = generateStatement(funcSymbol->consts.function.functionNamespace, tree);

			ret->insert(ret->end(), stat->begin(), stat->end());
			delete stat;
		}

		return ret;
	}

	ArgList* IMLGenerator::genFunctions(){
		ArgList* ret = new ArgList;
		_symbolMap::iterator it;
		for(it = globalNamespace.FunctionMap.begin(); it != globalNamespace.FunctionMap.end(); it++){
			Symbol* sym = it->second;
			ArgList* func = generateFunction(sym);
			// Create a new block.
			Quad* newBlock = new Quad;
			newBlock->op = BLOCK;
			newBlock->inner.blockChildren = new ArgList;
			newBlock->inner.blockChildren->insert(newBlock->inner.blockChildren->end(),func->begin(), func->end());
			ret->push_back(newBlock);
		}

		return ret;
	}

	void IMLGenerator::setupFunctionParameters(Symbol* funcSymbol){
		SymbolList* paramList = funcSymbol->consts.function.paramDeclaration;
		int currentParamOffset = 8;	// Skip over EBP storage and return address
		SymbolList::iterator it;
		for( it = paramList->begin(); it!=paramList->end(); it++){
			Symbol* sym = *it;

			sym->location = new Arg;
			sym->location->arg = EBP_OFFSET;
			sym->location->name.intName = currentParamOffset;

			// Increase offset by the size of the parameter.
			currentParamOffset += sym->type.size;
		}
	}

	int IMLGenerator::setupFunctionLocals(Symbol* funcSymbol){
		Namespace* funcSpace = funcSymbol->consts.function.functionNamespace;
		_symbolMap::iterator it;
		int currentLocalOffset = 0;
		for( it = funcSpace->SymbolMap.begin(); it!= funcSpace->SymbolMap.end(); it++){
			Symbol* sym = it->second;

			if(sym->type.primitiveType == PRIM_TYPE_POINTER)
				sym->type.size = 4;
			sym->location = new Arg;
			if( sym->scope == SCOPE_LOCAL ){
				// Decrease offset by the size of the element.
				currentLocalOffset -= sym->type.size;
				
				sym->location->arg = STACK;
				sym->location->name.intName = currentLocalOffset;
				sym->locations = new LocationList;
				sym->locations->push_back(sym->location);
			}
		}

		SymbolList::iterator it2;
		int currentParamOffset = 8;
		// Set up the parameters.
		for( it2 = funcSymbol->consts.function.paramDeclaration->begin(); it2 != funcSymbol->consts.function.paramDeclaration->end(); it2++){
			Symbol* sym = *it2;

			sym->location = new Arg;
			sym->location->arg = STACK;
			sym->location->name.intName = currentParamOffset;
			sym->locations = new LocationList;
			sym->locations->push_back(sym->location);

			currentParamOffset += sym->type.size;
		}
		return currentLocalOffset;
	}

	ArgList* IMLGenerator::generateFunction(Symbol* funcSymbol){
		// Set up the parameters.
		setupFunctionParameters(funcSymbol);

		ArgList* ret = new ArgList;
		// Create the function declaration
		GlobalDeclaration decl;
		decl.globalType = LABEL;
		decl.symbolName = &funcSymbol->symbolName;
		Quad* quad = new Quad;
		quad->op = GLOBAL_DECL;
		quad->inner.declaration = decl;
		ret->push_back(quad);

		// Push ebp
		Quad* pushEbp = new Quad;
		pushEbp->op = PUSH;
		Arg ebp;
		ebp.arg = EBP;
		pushEbp->inner.normalOp.arg1 = ebp;
		ret->push_back(pushEbp);

		// Make ebp the current stack pointer ( after the old EBP )
		quad = new Quad;
		quad->op =MOV;
		quad->inner.normalOp.arg1.arg = ESP;
		quad->inner.normalOp.arg2.arg = EBP;
		ret->push_back(quad);

		// Set up the locals.
		int totalLocalSize = setupFunctionLocals(funcSymbol);
		// Create space for the locals
		quad = new Quad;
		quad->op = SUB;
		quad->inner.normalOp.arg1.arg = INT_CONSTANT;
		quad->inner.normalOp.arg1.name.intName = -totalLocalSize;	// Remove the negative sign
		quad->inner.normalOp.arg2.arg = ESP;
		ret->push_back(quad);

		// Do the function itself.
		ArgList* firstPassBody = generateFunctionBody(funcSymbol);
		// Do the second pass
		ArgList* body = doSecondPass(funcSymbol->consts.function.functionNamespace,firstPassBody);
		delete firstPassBody;

		ret->insert(ret->end(), body->begin(), body->end());
		delete body;

		// Clear locals.
		quad = new Quad;
		quad->op = MOV;
		quad->inner.normalOp.arg1.arg = EBP;
		quad->inner.normalOp.arg2.arg = ESP;
		ret->push_back(quad);
		// Restore EBP
		quad = new Quad;
		quad->op = POP;
		quad->inner.normalOp.arg1.arg = EBP;
		ret->push_back(quad);
		// Return;
		quad = new Quad;
		quad->op = RETURN;
		ret->push_back(quad);

		return ret;
	}

	ArgList* IMLGenerator::genGlobals(_symbolMap& symbolMap){
		ArgList* ret = new ArgList;

		// Generate the constants first.
		if( &symbolMap == &globalNamespace.SymbolMap){
			ArgList* consts = genGlobals(globalNamespace.ConstantMap);
			ret->insert(ret->end(), consts->begin(), consts->end());
			delete consts;
		}

		_symbolMap::iterator it;
		for( it = symbolMap.begin(); it != symbolMap.end(); it++){
			Symbol* sym = it->second;
			sym->location = new Arg;
			sym->location->arg = GLOBAL;
			sym->location->name.symbolName = &sym->symbolName;
			sym->locations = new LocationList;
			sym->locations->push_back(sym->location);
			Quad* newLine = new Quad;
			newLine->op = GLOBAL_DECL;

			// Alias declaration
			GlobalDeclaration& decl = newLine->inner.declaration;

			decl.symbolName = & sym->symbolName;

			// Assuming he has a value
		
			switch(sym->type.primitiveType){
			case PRIM_TYPE_INT:
				decl.globalType = INT;
				decl.value.intValue = sym->consts.exprConst->inner.nodeToken.value.intValue;
				break;
			case PRIM_TYPE_FLOAT:
				decl.globalType = FLOAT;
				decl.value.floatValue = sym->consts.exprConst->inner.nodeToken.value.floatValue;
				break;
			case PRIM_TYPE_ARRAY:
			case PRIM_TYPE_STRING:
				decl = createArrayDeclaration(sym);
				break;
			case PRIM_TYPE_POINTER:
				decl.globalType = POINTER; // Pointers are just arrays.
				decl.value.arrayValue = new vector<GlobalDeclaration>;
				GlobalDeclaration pointTo;
				pointTo.globalType = SYMBOL;
				pointTo.symbolName = &sym->consts.exprConst->inner.sym->symbolName;
				decl.value.arrayValue->push_back(pointTo);
				break;

			}

			ret->push_back(newLine);
		}

		return ret;
	}


	void IMLGenerator::generateIML(){
		identation = "";
		cout << "Globals:" << endl;
		identation = "\t";
		ArgList* globals = genGlobals(globalNamespace.SymbolMap);
		printArgList(globals, cout);

		// Functions
		cout << "\nText:\n" <<endl;
		identation = "";
		ArgList* functions = genFunctions();
		printArgList(functions, cout);

		int i = 0;
	}


}
